import torch

def _make_grid(nx=20, ny=20):
    yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
    print('torch.arange(ny) is ', torch.arange(ny))
    print('torch.arange(nx) is ', torch.arange(nx))
    print(yv.size(), xv.size())
    print((torch.stack((xv, yv), 2)).size())
    print(torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).size())
    return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()

_make_grid()

n, gd = 0.5, 0.67
n = max(round(n * gd), 1) if n > 1 else n  # depth gain
print(n)